from PyQt6.QtCore import QThread, pyqtSignal
from PyQt6.QtWidgets import (QApplication, QMainWindow, QPushButton, QDialog, QVBoxLayout, QListWidget, QLabel, QComboBox, QLineEdit, QHBoxLayout, QFileDialog, QProgressBar, QMessageBox)
import sys
import os
import sqlite3
import struct
import psutil
import wmi
import ctypes
from ctypes import wintypes

class DiskDialog(QDialog):
    def __init__(self, disks):
        super().__init__()
        self.setWindowTitle("选择磁盘")
        self.setFixedSize(300, 200)
        self.layout = QVBoxLayout()
        self.list_widget = QListWidget()
        self.list_widget.addItems(disks)
        self.layout.addWidget(self.list_widget)
        self.ok_button = QPushButton("确定")
        self.ok_button.clicked.connect(self.accept)
        self.layout.addWidget(self.ok_button)
        self.setLayout(self.layout)

    def selected_disk(self):
        return self.list_widget.currentItem().text()

def get_disks():
    disks = []
    # 获取逻辑磁盘（如 C:\、D:\）
    for partition in psutil.disk_partitions():
        # if 'fixed' in partition.opts:  # 只获取固定磁盘
        disks.append('\\\\.\\'+partition.device[0:-1])
    # 创建 WMI 客户端
    c = wmi.WMI()
    # 获取所有物理磁盘
    physical_disks = c.query("SELECT * FROM Win32_DiskDrive")
    for disk in physical_disks:
        disks.append(disk.DeviceID)
    return disks

def get_physical_drive_size(drive_path):
    # 定义常量
    IOCTL_DISK_GET_LENGTH_INFO = 0x7405C
    GENERIC_READ = 0x80000000
    FILE_SHARE_READ = 0x1
    OPEN_EXISTING = 0x3
    # 定义结构体
    class LARGE_INTEGER(ctypes.Union):
        class Struct1(ctypes.Structure):
            _fields_ = [("LowPart", wintypes.DWORD),
                        ("HighPart", wintypes.LONG)]
        _fields_ = [("li", Struct1),
                    ("QuadPart", ctypes.c_longlong)]
    # 打开物理磁盘
    handle = ctypes.windll.kernel32.CreateFileW(
        drive_path,
        GENERIC_READ,
        FILE_SHARE_READ,
        None,
        OPEN_EXISTING,
        0,
        None
    )
    if handle == wintypes.HANDLE(-1).value:
        raise ctypes.WinError(ctypes.get_last_error())
    size = LARGE_INTEGER()
    bytes_returned = wintypes.DWORD(0)
    try:
        # 获取磁盘大小
        result = ctypes.windll.kernel32.DeviceIoControl(
            handle,
            IOCTL_DISK_GET_LENGTH_INFO,
            None,
            0,
            ctypes.byref(size),
            ctypes.sizeof(size),
            ctypes.byref(bytes_returned),
            None
        )
        if not result:
            raise ctypes.WinError(ctypes.get_last_error())
        return size.QuadPart
    finally:
        ctypes.windll.kernel32.CloseHandle(handle)

def get_disk_size(disk_path):
    if disk_path.startswith("\\\\.\\"):  # 物理磁盘
        try:
            size = get_physical_drive_size(disk_path)
            return size
        except OSError:
            return 0
    else:  # 逻辑磁盘
        return os.path.getsize(disk_path)

class ScanThread(QThread):
    progress_signal = pyqtSignal(int, int, float, float)  # 扫描进度信号
    export_progress_signal = pyqtSignal(int, int, float)  # 导出进度信号

    def __init__(self, img_path, exp_dir, db_path):
        super().__init__()
        self.img_path = img_path
        self.exp_dir = exp_dir
        self.db_path = db_path

    def run(self):
        db = sqlite3.connect(self.db_path)
        cursor = db.cursor()
        try:
            img = open(self.img_path, 'rb', buffering=1048576)
        except FileNotFoundError:
            print(f"无法打开磁盘: {self.img_path}")
            return
        except PermissionError:
            print(f"权限不足，无法打开磁盘: {self.img_path}")
            return
        img_size = get_disk_size(self.img_path)
        if img_size == 0:
            print(f"无法获取磁盘大小: {self.img_path}")
            return
        cursor.execute('create table if not exists info(space_id,page_id,indx_id,offset,img_path)')
        n = 0
        page_num = 0
        while n * 512 < img_size:
            try:
                img.seek(n * 512)
                block = img.read(16384)
                top_check1 = int.from_bytes(block[:4], 'big')
                top_check2 = int.from_bytes(block[20:24], 'big')
                tail_check1 = int.from_bytes(block[-8:-4], 'big')
                tail_check2 = int.from_bytes(block[-4:], 'big')
                if top_check2 == tail_check2 and top_check1 != 0 and top_check2 != 0 and tail_check1 != 0:
                    pre_page = struct.unpack('>L', block[0x08:0x0c])[0]
                    next_page = struct.unpack('>L', block[0x0c:0x10])[0]
                    page_type = struct.unpack('>H', block[0x18:0x1A])[0]
                    lsn = struct.unpack('>Q', block[0x1A:0x22])[0]
                    f_off = struct.unpack('>H', block[16374:16376])[0]
                    if f_off < 512 and f_off > 0 and page_type != 0x0a:
                        magic = block[f_off:f_off + 7].decode(encoding='utf-8', errors='ignore')
                        if magic == 'infimum':
                            page_id = struct.unpack('>L', block[4:8])[0]
                            space_id = struct.unpack('>L', block[34:38])[0]
                            indx_id = struct.unpack('>L', block[0x46:0x4a])[0]
                            if space_id != 4278124286:
                                cursor.execute(
                                    rf'insert into info values({space_id},{page_id},{indx_id},{n * 512},"{self.img_path}")')
                            n += 32
                            page_num += 1
                        else:
                            n += 1
                    elif (page_type == 0x000a or page_type == 0x0017 or page_type == 0x0018) and lsn == 0:
                        page_id = struct.unpack('>L', block[4:8])[0]
                        space_id = struct.unpack('>L', block[34:38])[0]
                        indx_id = 999999
                        cursor.execute(
                            rf'insert into info values({space_id},{page_id},{indx_id},{n * 512},"{self.img_path}")')
                        n += 32
                        page_num += 1
                    else:
                        n += 1
                else:
                    n += 1
                if n % 100000 == 0:
                    self.progress_signal.emit(n * 512, img_size, (n * 512 / img_size) * 100, page_num * 16384 / 1024 / 1024)
                    db.commit()
            except:
                break
        db.commit()
        cursor.execute(r'create index if not exists page_idx on info(page_id)')
        cursor.execute(r'create index if not exists space_idx on info(space_id)')
        cursor.execute(r'create index if not exists indx_idx on info(indx_id)')
        # 导出操作
        record = cursor.execute('select * from info group by space_id').fetchall()
        total_records = len(record)
        for i, rcd in enumerate(record):
            name = str(rcd[0])
            if int(name) == 0:
                for indx_id in (1, 3):
                    db_file_path = os.path.join(self.exp_dir, f'{name}_{indx_id}')
                    db_file = open(db_file_path, 'wb')
                    page_record = cursor.execute(
                        f'select * from info where space_id = 0 and indx_id = {indx_id} order by page_id').fetchall()
                    for page_rcd in page_record:
                        img_path = page_rcd[-1]
                        img = open(img_path, 'rb')
                        page_id = page_rcd[1]
                        offset = page_rcd[3]
                        img.seek(offset)
                        page = img.read(16384)
                        db_file.write(page)
                db_file_path = os.path.join(self.exp_dir, f'{name}')
                db_file = open(db_file_path, 'wb')
                page_record = cursor.execute(
                    f'select * from info where space_id = 0 order by page_id').fetchall()
                for page_rcd in page_record:
                    img_path = page_rcd[-1]
                    img = open(img_path, 'rb')
                    page_id = page_rcd[1]
                    offset = page_rcd[3]
                    img.seek(offset)
                    page = img.read(16384)
                    db_file.write(page)
            else:
                db_file_path = os.path.join(self.exp_dir, name)
                db_file = open(db_file_path, 'wb')
                page_record = cursor.execute(f'select * from info where space_id = {rcd[0]} order by page_id').fetchall()
                for page_rcd in page_record:
                    img_path = page_rcd[-1]
                    img = open(img_path, 'rb')
                    page_id = page_rcd[1]
                    offset = page_rcd[3]
                    img.seek(offset)
                    page = img.read(16384)
                    db_file.seek(page_id * 16384)
                    db_file.write(page)
            self.export_progress_signal.emit(i + 1, total_records, ((i + 1) / total_records) * 100)
        db.commit()
        db.close()

class ScanProgressDialog(QDialog):
    def __init__(self,parent=None):
        super().__init__(parent)
        self.setWindowTitle("进度")
        self.setStyleSheet("""
                              QDialog { background-color: #f0f0f0; }
                              QPushButton { background-color: #c1c1c1; color: black; border-radius: 5px; padding: 5px 15px; }
                              QPushButton:hover { background-color: #0053a4; }
                              QLineEdit { border: 1px solid #ccc; border-radius: 5px; padding: 5px; }
                              QPushButton:disabled { background-color: #e0e0e0; color: #a0a0a0; }  # 禁用状态的样式   
                              QLabel { color: #333; }
                              QLineEdit::placeholder { color: #aaa; }
                          """)
        self.setFixedSize(400, 200)
        self.layout = QVBoxLayout()
        # 扫描进度条
        self.scan_progress_label = QLabel("扫描进度:")
        .0
        self.scan_progress_bar = QProgressBar()
        self.layout.addWidget(self.scan_progress_label)
        self.layout.addWidget(self.scan_progress_bar)
        # 导出进度条
        self.export_progress_label = QLabel("导出进度:")
        self.export_progress_bar = QProgressBar()
        self.layout.addWidget(self.export_progress_label)
        self.layout.addWidget(self.export_progress_bar)
        self.setLayout(self.layout)

    def update_scan_progress(self, current, total, percent, size):
        self.scan_progress_bar.setValue(int(percent))
        self.scan_progress_bar.setFormat(f"扫描进度: {percent:.2f}% 已扫描到数据页: {size:.2f}M")

    def update_export_progress(self, current, total, percent):
        self.export_progress_bar.setValue(int(percent))
        self.export_progress_bar.setFormat(f"导出进度: {percent:.2f}% ({current}/{total})")

class ScanDialog(QDialog):
    def __init__(self,parent=None):
        super().__init__(parent)
        self.setWindowTitle("扫描设置")
        self.setStyleSheet("""
                      QDialog { background-color: #f0f0f0; }
                      QPushButton { background-color: #c1c1c1; color: black; border-radius: 5px; padding: 5px 15px; }
                      QPushButton:hover { background-color: #0053a4; }
                      QLineEdit { border: 1px solid #ccc; border-radius: 5px; padding: 5px; }
                      QPushButton:disabled { background-color: #e0e0e0; color: #a0a0a0; }  # 禁用状态的样式   
                      QLabel { color: #333; }
                      QLineEdit::placeholder { color: #aaa; }
                  """)
        self.setFixedSize(600, 300)
        self.layout = QVBoxLayout()
        # 扫描类型选择
        self.scan_type_label = QLabel("选择扫描类型:")
        self.scan_type_combo = QComboBox()
        self.scan_type_combo.addItems(["镜像文件", "硬盘"])
        self.layout.addWidget(self.scan_type_label)
        self.layout.addWidget(self.scan_type_combo)
        # 镜像文件或硬盘选择
        self.img_path_label = QLabel("选择镜像文件或硬盘:")
        self.img_path_layout = QHBoxLayout()
        self.img_path_edit = QLineEdit()
        self.img_path_button = QPushButton("浏览")
        self.img_path_button.setFixedWidth(80)
        self.img_path_button.clicked.connect(self.select_img_path)
        self.img_path_layout.addWidget(self.img_path_edit)
        self.img_path_layout.addWidget(self.img_path_button)
        self.layout.addWidget(self.img_path_label)
        self.layout.addLayout(self.img_path_layout)
        # 导出目录选择
        self.exp_dir_label = QLabel("选择导出目录:")
        self.exp_dir_layout = QHBoxLayout()
        self.exp_dir_edit = QLineEdit()
        self.exp_dir_button = QPushButton("浏览")
        self.exp_dir_button.setFixedWidth(80)
        self.exp_dir_button.clicked.connect(self.select_exp_dir)
        self.exp_dir_layout.addWidget(self.exp_dir_edit)
        self.exp_dir_layout.addWidget(self.exp_dir_button)
        self.layout.addWidget(self.exp_dir_label)
        self.layout.addLayout(self.exp_dir_layout)
        # 数据库路径选择
        self.db_path_label = QLabel("选择数据库路径:")
        self.db_path_layout = QHBoxLayout()
        self.db_path_edit = QLineEdit()
        self.db_path_button = QPushButton("浏览")
        self.db_path_button.setFixedWidth(80)
        self.db_path_button.clicked.connect(self.select_db_path)
        self.db_path_layout.addWidget(self.db_path_edit)
        self.db_path_layout.addWidget(self.db_path_button)
        self.layout.addWidget(self.db_path_label)
        self.layout.addLayout(self.db_path_layout)
        # 确定按钮
        self.ok_button = QPushButton("确定")
        self.ok_button.clicked.connect(self.accept)
        self.layout.addWidget(self.ok_button)
        self.setLayout(self.layout)

    def select_img_path(self):
        if self.scan_type_combo.currentText() == "镜像文件":
            file_path, _ = QFileDialog.getOpenFileName(self, "选择镜像文件", "", "镜像文件(*.img *.iso *.vmdk);;所有文件 (*)")
            if file_path:
                self.img_path_edit.setText(file_path)
        else:
            # 显示逻辑磁盘和物理磁盘
            disks = get_disks()
            dialog = DiskDialog(disks)
            if dialog.exec():
                selected_disk = dialog.selected_disk()
                self.img_path_edit.setText(selected_disk)

    def select_exp_dir(self):
        dir_path = QFileDialog.getExistingDirectory(self, "选择导出目录")
        if dir_path:
            self.exp_dir_edit.setText(dir_path)

    def select_db_path(self):
        file_path, _ = QFileDialog.getSaveFileName(self, "选择SQLite数据库路径", "", "SQLite Files (*.db)")
        if file_path:
            self.db_path_edit.setText(file_path)